from threading import local
import torch
from transformers import AutoModel, AutoProcessor
from transformers.image_utils import load_image
import numpy as np
from fastapi import FastAPI, UploadFile, File, Form, HTTPException,Body
from fastapi.responses import JSONResponse
import io
from io import BytesIO
from PIL import Image
import uvicorn
import requests
import os 

# 设置离线模式，如需在线下载模型请注释这一行
os.environ["TRANSFORMERS_OFFLINE"] = "1" 

app = FastAPI(title="SigLIP2图像特征提取API")

# 使用siglip2-so400m-patch14-384模型
#model_name = "google/siglip2-so400m-patch14-384"
model_name = "OFA-Sys/chinese-clip-vit-base-patch16"
# 加载模型和处理器
print(f"正在加载模型: {model_name}")
model = AutoModel.from_pretrained(model_name, torch_dtype=torch.float32).eval()
processor = AutoProcessor.from_pretrained(model_name)
print(f"模型加载成功，词汇表大小: {model.config.text_config.vocab_size}")


# 提取图像特征的函数
def extract_image_features(image):
    # 预处理图像
    inputs = processor(images=image, return_tensors="pt").to(model.device)
    # 前向传播，获取图像特征
    with torch.no_grad():
        image_features = model.get_image_features(**inputs)
    # 归一化特征
    image_embeds = image_features / image_features.norm(dim=-1, keepdim=True)
    # 转换为numpy数组并转为列表返回
    return image_embeds.cpu().numpy().astype(np.float32)[0].tolist()

@app.get("/text")
def get_text_features(text: str):
    with torch.no_grad():
        text_inputs = processor(text=[text],  return_tensors="pt", padding=True)
        text_features = model.get_text_features(**text_inputs)
        text_embeds = text_features / text_features.norm(dim=-1, keepdim=True)
    features = text_embeds.cpu().numpy().astype(np.float32)[0].tolist()
    return {"data": features}

@app.get("/image")
async def get_image_features(url: str):
    try:
        data = requests.get(url).content
        image = Image.open(io.BytesIO(data)).convert("RGB")
        features = extract_image_features(image)
        return {"data": features}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理图像时出错: {str(e)}")

@app.post("/image/binary")
async def get_features_by_stream(data: bytes = Body(...)):
    """
    通过字节流提取图像特征
    """
    try:
        # 直接使用字节数据创建图像
        image = Image.open(io.BytesIO(data)).convert("RGB")
        features = extract_image_features(image)
        return {"data": features}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理图像时出错: {str(e)}")

@app.get("/model_info")
async def get_model_info():
    print(processor.tokenizer)
    return {
        "model_type": model.config.model_type,
        "vocab_size": model.config.text_config.vocab_size,
        "supports_chinese": "中文" in processor.tokenizer.vocab or any(chr(i) in processor.tokenizer.vocab for i in range(19968, 40959))  # 基本汉字范围
    }

if __name__ == "__main__":
    # 启动服务器
    uvicorn.run(app, host="0.0.0.0", port=8060)